package jtg.generator;

import com.microsoft.z3.*;
import fj.Hash;
import jtg.graphics.SootCFG;
import jtg.parser.Z3Parser;
import jtg.state.PathState;
import jtg.trie.TrieNode;
import jtg.trie.TrieTree;
import jtg.visualizer.Visualizer;
import soot.Body;
import soot.Local;
import soot.Unit;
import soot.Value;
import soot.jimple.internal.*;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

public class PrimeCoverageGenerator extends Generator {


    public PrimeCoverageGenerator(String className, String methodName) {
        super(className, methodName);
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new PrimeCoverageGenerator(defaultClsPath, className, methodName);
    }

    public PrimeCoverageGenerator(String classPath, String className, String methodName) {
        super(classPath, className, methodName);
    }

    @Override
    void drawCFG() {
        drawCFG(clsName + "." + mtdName, false);
    }

    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }


    private List<Local> getJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<Local>();
        for (Local l : body.getLocals()) {
            if (l.toString().startsWith("$")) jimpleVars.add(l);
        }
        return jimpleVars;
    }

    private List<Local> getParameters() {
        return new ArrayList<>(body.getParameterLocals());
    }

    public List<String> generate() {

        List<Unit> path = null;
        ArrayList<String> testSet = null;

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        try {
            testSet = new ArrayList<String>();
            Context context = new Context();
            Solver solver = context.mkSolver();

            if(ug.getHeads().size()>1) throw new Exception("The number of heads is more than 1");
            Unit h = ug.getHeads().get(0);
            TrieTree primePathTree = getPrimePathTree(ug);
            // path = ug.getExtendedBasicBlockPathBetween(h, t);
            List<PathState> pathStates = initPathState(context, h);
            // 广搜开始,要在其中记录路径约束
            List<Unit> visitedNode = new ArrayList<>();
            calculateWithBFS(context, ug, primePathTree, visitedNode, pathStates, testSet);
            // 计算覆盖率，遍历字典树
            int primePathCount = 0;
            int visitedPrimePathCount = 0;
            int[] t = primePathTree.calTreeNode();
            List<List<Unit>> primePath = new ArrayList<>();
            primePath = primePathTree.getAllBranch(new ArrayList<Unit>(), primePathTree.root, new ArrayList<>());
            System.out.println("The primePath is: ");
            for(List<Unit> p : primePath){
                System.out.println(p);
            }
            primePathCount = t[0];
            visitedPrimePathCount = t[1];
            System.out.println("The primePathCount is: " + primePathCount+ " The visitedPrimePathCount is: " + visitedPrimePathCount);
        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
            e.printStackTrace();
        }
        if (!testSet.isEmpty()) {
            System.out.println("");
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) " + tc.toString());
                count++;
            }
        }
        return testSet;
    }

    public List<PathState> initPathState(Context context, Unit h){
        List<PathState> pathStates = new ArrayList<>();
        // 创建初始状态
        PathState initialState = new PathState(h);
        for(Local local : getParameters()){
            String parName = local.getName();
            initialState.params.put(parName, context.mkIntConst(parName));
            initialState.paramLatestMap.put(parName,0);
        }
        initialState.unitPath.add(h);
        pathStates.add(initialState);
        return pathStates;
    }

    public void calculateWithBFS(Context context, UnitGraph ug, TrieTree primePathTree, List<Unit> visitedNode, List<PathState> pathStates, List<String> testSet) throws Exception {
        while(!pathStates.isEmpty()){

            PathState pathState = pathStates.remove(0);
            if(pathState.unitPath.size()>100) return;
            Unit unit = pathState.node;
            visitedNode.add(unit);

            Solver solver = context.mkSolver();
            List<BoolExpr> pathConstraint = calPathConstraint(context, pathState.unitPath);
            if(!pathConstraint.isEmpty()) {
                for(BoolExpr expr : pathConstraint){
                    solver.add(expr);
                }

                if(solver.check() == Status.UNSATISFIABLE){
                    continue;
                }
            }

            Collections.reverse(pathState.unitPath);
            TrieNode tNode = primePathTree.startWith(pathState.unitPath);
            Collections.reverse(pathState.unitPath);
            if(tNode != null){
//                System.out.println("The prime path is: " + pathState.unitPath);
                String test = getTestSetByPathConstraint(context, pathConstraint);
                if(test!=null){
                    tNode.isVisited = true;
                    testSet.add(test);
                }
            }

            if(unit instanceof JReturnStmt||unit instanceof JReturnVoidStmt){
                // 如果是return语句，那么就是一条路径
                // 生成一个新的状态
//                if(pathState.unitPath.size()<=15){
//                    System.out.println("return path:"+pathState.unitPath);
//                }
            }else if(unit instanceof JIfStmt){
                // 如果是if语句，那么就是两条路径
                // 生成两个新的状态
                PathState pathState1 = new PathState(pathState);
                PathState pathState2 = new PathState(pathState);
                // 生成两个新的约束
                // 有可能是一个循环，判断一下

                // BoolExpr ifstms = (BoolExpr) getZ3ExprFromSootExpr(context, pathState.params, ((JIfStmt) unit).getCondition());
                List<Unit> nextUnitList = ug.getSuccsOf(unit);
                Unit nextUnit1 = nextUnitList.get(0);
                Unit nextUnit2 = nextUnitList.get(1);
                pathState1.node = nextUnit1;
                pathState2.node = nextUnit2;
                pathState1.unitPath.add(nextUnit1);
                pathState2.unitPath.add(nextUnit2);
                // pathState1.PathConstraints.add(context.mkNot(ifstms));
                // pathState2.PathConstraints.add(ifstms);


                pathStates.add(pathState1);
                pathStates.add(pathState2);
            }else if(unit instanceof JAssignStmt){
                // 如果是赋值语句，那么就是一条路径
                // 生成一个新的状态
                PathState pathState1 = new PathState(pathState);
                // 生成一个新的约束
                String leftOp = ((JAssignStmt) unit).getLeftOp().toString();
                pathState1.params.put(leftOp, context.mkIntConst(leftOp));
                // Expr rightOp = getZ3ExprFromSootExpr(context, pathState1.params, ((JAssignStmt) unit).getRightOp());
                // if (rightOp == null) {
                //    System.out.println("Ignore the statement: " + unit.toString());
                //    continue;
                //}
                // pathState1.PathConstraints.add(context.mkEq(pathState1.params.get(leftOp), rightOp));
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                pathStates.add(pathState1);

            }else if(unit instanceof JGotoStmt){
                // 如果是goto语句，那么就是一条路径
                // 生成一个新的状态
                PathState pathState1 = new PathState(pathState);
                // 判断是否是循环
                Unit nextUnit = ug.getSuccsOf(unit).get(0);

                pathState1.node = nextUnit;
                pathState1.unitPath.add(pathState1.node);

                pathStates.add(pathState1);
            }else {
                // 如果是其它语句，那么就是一条路径
                // 生成一个新的状态
                PathState pathState1 = new PathState(pathState);
                // 生成一个新的约束
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                TrieNode trieNode = primePathTree.search(pathState1.unitPath);
                pathStates.add(pathState1);
            }
        }
    }


    // 深搜找所有的基路径，并反着加入字典树
    public TrieTree getPrimePathTree(UnitGraph ug){
        TrieTree primePathTree = new TrieTree();
        for(Unit unit: ug){
            List<Unit> path = new ArrayList<>();
            List<Unit> visited = new ArrayList<>();
            primePathDFSHelper(ug, unit,unit, primePathTree, path, visited);
        }
//        System.out.println("The primePathTree is: " + primePathTree);
        return primePathTree;
    }

    private void primePathDFSHelper(UnitGraph ug, Unit unit, Unit head, TrieTree primePathTree, List<Unit> path, List<Unit> visited){
        path.add(unit);
        boolean isEnd = visited.contains(unit);
        visited.add(unit);
        if(unit instanceof JReturnStmt || unit instanceof JReturnVoidStmt){
            // 如果是return语句，那么就是一条路径
            // 判断这个路径的末端是否是基路径
            isEnd = true;
        }else {
            List<Unit> nextUnitList = ug.getSuccsOf(unit);
            for(Unit nextUnit : nextUnitList) {
                if(visited.contains(nextUnit)){
                    if(nextUnit.equals(head)) {

                        path.add(nextUnit);
                        // 倒着加入字典树
//                        System.out.println(path);
                        Collections.reverse(path);
                        primePathTree.insert(path);
                        Collections.reverse(path);
                        path.remove(path.size()-1);
                    }else{
//                        System.out.println("1" + path);
                        // 倒着加入字典树
                        Collections.reverse(path);
                        primePathTree.insert(path);
                        Collections.reverse(path);
                    }

                }else primePathDFSHelper(ug, nextUnit, head, primePathTree, path, visited);
            }
        }
//        System.out.println(isEnd);
        if(isEnd){
            // 如果是基路径，那么就加入字典树
//            System.out.println("2"+path);
            Collections.reverse(path);
            primePathTree.insert(path);
            Collections.reverse(path);
        }
        path.remove(unit);
        visited.remove(unit);
    }

}
